package com.googlecode.totallylazy.lambda.weaving;

import com.googlecode.totallylazy.lambda.weaving.LambdaTreeWeaver.MethodAnalyzer.LambdaAnalyzer;
import org.objectweb.asm.ClassReader;
import org.objectweb.asm.ClassWriter;
import org.objectweb.asm.Opcodes;
import org.objectweb.asm.tree.ClassNode;

import java.util.HashMap;
import java.util.Map;

import static com.googlecode.totallylazy.lambda.weaving.Debug.debug;
import static org.objectweb.asm.ClassWriter.COMPUTE_FRAMES;
import static org.objectweb.asm.Type.getObjectType;

public class LambdaTreeTransformer implements Opcodes {
    Map<String, byte[]> lambdasByClassName = new HashMap<String, byte[]>();

    public LambdaTreeTransformer() {
        debug("current class loader is " + getClass().getClassLoader());
    }

    public Map<String, byte[]> getLambdasByClassName() {
        return lambdasByClassName;
    }

    public byte[] transform(ClassLoader loader, ClassFilter filter, String name, byte[] in) throws Exception {
        if (lambdasByClassName.containsKey(name)) {
            debug("generated lambda requested by the class loader " + name);
            return lambdasByClassName.get(name);
        }

        ClassReader cr = new ClassReader(in);

        LambdaTreeWeaver weaver = new LambdaTreeWeaver(loader, filter, cr);
        ClassNode classNode = weaver.analyze().transform();

        if (!weaver.hasLambdas())
            return null;

        ClassWriter classWriter = new ClassWriter(COMPUTE_FRAMES);
        classNode.accept(classWriter);

        byte[] bs = classWriter.toByteArray();

        for (LambdaAnalyzer lambdaAnalyzer : weaver.getLambdas()) {
            classWriter = new ClassWriter(COMPUTE_FRAMES);
            lambdaAnalyzer.lambda.accept(classWriter);
            newLambdaClass(getObjectType(lambdaAnalyzer.lambdaClass()).getClassName(), classWriter.toByteArray());
        }

        return bs;
    }

    void newLambdaClass(String name, byte[] bs) {
        lambdasByClassName.put(name, bs);

    }
}
